import pymysql
import pandas as pd
import logging
from typing import List, Dict, Any, Optional, Union
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from app.config import settings

logger = logging.getLogger(__name__)


class DorisClient:
    """Doris 数据库客户端"""
    
    def __init__(self, 
                 host: str = None, 
                 port: int = None, 
                 database: str = None,
                 username: str = None, 
                 password: str = None):
        """
        初始化 Doris 客户端
        
        Args:
            host: Doris FE 地址
            port: Doris 查询端口
            database: 数据库名
            username: 用户名
            password: 密码
        """
        self.host = host or settings.DORIS_HOST
        self.port = port or settings.DORIS_QUERY_PORT
        self.database = database or settings.DORIS_DATABASE
        self.username = username or settings.DORIS_USERNAME
        self.password = password or settings.DORIS_PASSWORD
        
        # 创建连接URL
        self.connection_url = f"mysql+pymysql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
        
        # SQLAlchemy 引擎
        self._engine: Optional[Engine] = None
    
    @property
    def engine(self) -> Engine:
        """获取 SQLAlchemy 引擎"""
        if self._engine is None:
            self._engine = create_engine(
                self.connection_url,
                pool_pre_ping=True,
                pool_recycle=3600,
                echo=settings.DEBUG
            )
        return self._engine
    
    def test_connection(self) -> bool:
        """测试数据库连接"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text("SELECT 1 as test"))
                return result.fetchone()[0] == 1
        except Exception as e:
            logger.error(f"Doris 连接测试失败: {str(e)}")
            return False
    
    def execute_query(self, sql: str, params: Dict[str, Any] = None) -> List[Dict[str, Any]]:
        """
        执行查询语句
        
        Args:
            sql: SQL 查询语句
            params: 查询参数
            
        Returns:
            查询结果列表
        """
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), params or {})
                columns = result.keys()
                rows = result.fetchall()
                return [dict(zip(columns, row)) for row in rows]
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            raise
    
    def execute_sql(self, sql: str, params: Dict[str, Any] = None) -> int:
        """
        执行SQL语句（INSERT、UPDATE、DELETE等）
        
        Args:
            sql: SQL 语句
            params: 参数
            
        Returns:
            影响的行数
        """
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), params or {})
                conn.commit()
                return result.rowcount
        except Exception as e:
            logger.error(f"执行SQL失败: {str(e)}")
            raise
    
    def get_tables(self) -> List[Dict[str, Any]]:
        """获取数据库中的所有表"""
        sql = """
        SELECT 
            TABLE_NAME as table_name,
            TABLE_COMMENT as table_comment,
            TABLE_TYPE as table_type
        FROM information_schema.TABLES 
        WHERE TABLE_SCHEMA = %s
        ORDER BY TABLE_NAME
        """
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), {"TABLE_SCHEMA": self.database})
                return [dict(row._mapping) for row in result]
        except Exception as e:
            logger.error(f"获取表列表失败: {str(e)}")
            return []
    
    def get_table_schema(self, table_name: str) -> List[Dict[str, Any]]:
        """获取表结构"""
        sql = """
        SELECT 
            COLUMN_NAME as column_name,
            DATA_TYPE as data_type,
            IS_NULLABLE as is_nullable,
            COLUMN_DEFAULT as default_value,
            COLUMN_COMMENT as column_comment,
            COLUMN_KEY as column_key
        FROM information_schema.COLUMNS 
        WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
        ORDER BY ORDINAL_POSITION
        """
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), {
                    "TABLE_SCHEMA": self.database,
                    "TABLE_NAME": table_name
                })
                return [dict(row._mapping) for row in result]
        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            return []
    
    def create_table_from_dataframe(self, df: pd.DataFrame, table_name: str, 
                                   if_exists: str = 'replace') -> bool:
        """
        从 DataFrame 创建表并导入数据
        
        Args:
            df: pandas DataFrame
            table_name: 目标表名
            if_exists: 表存在时的处理方式 ('fail', 'replace', 'append')
            
        Returns:
            是否成功
        """
        try:
            df.to_sql(
                name=table_name,
                con=self.engine,
                if_exists=if_exists,
                index=False,
                method='multi'
            )
            logger.info(f"成功创建表 {table_name} 并导入 {len(df)} 条数据")
            return True
        except Exception as e:
            logger.error(f"创建表失败: {str(e)}")
            return False
    
    def bulk_insert(self, table_name: str, data: List[Dict[str, Any]]) -> bool:
        """
        批量插入数据
        
        Args:
            table_name: 目标表名
            data: 数据列表
            
        Returns:
            是否成功
        """
        if not data:
            return True
            
        try:
            df = pd.DataFrame(data)
            df.to_sql(
                name=table_name,
                con=self.engine,
                if_exists='append',
                index=False,
                method='multi'
            )
            logger.info(f"成功向表 {table_name} 插入 {len(data)} 条数据")
            return True
        except Exception as e:
            logger.error(f"批量插入失败: {str(e)}")
            return False
    
    def get_table_stats(self, table_name: str) -> Dict[str, Any]:
        """获取表统计信息"""
        try:
            # 获取行数
            count_sql = f"SELECT COUNT(*) as row_count FROM {table_name}"
            count_result = self.execute_query(count_sql)
            row_count = count_result[0]['row_count'] if count_result else 0
            
            # 获取表大小（这里简化处理，实际可以通过系统表获取）
            return {
                "row_count": row_count,
                "table_name": table_name,
                "database": self.database
            }
        except Exception as e:
            logger.error(f"获取表统计信息失败: {str(e)}")
            return {}
    
    def optimize_table(self, table_name: str) -> bool:
        """优化表（Doris 特有的优化操作）"""
        try:
            # Doris 的表优化可以通过 ALTER TABLE 语句
            sql = f"ALTER TABLE {table_name} COMPACT"
            self.execute_sql(sql)
            logger.info(f"成功优化表 {table_name}")
            return True
        except Exception as e:
            logger.error(f"优化表失败: {str(e)}")
            return False
    
    def close(self):
        """关闭连接"""
        if self._engine:
            self._engine.dispose()
            self._engine = None


# 全局 Doris 客户端实例
doris_client = DorisClient()


def get_doris_client() -> DorisClient:
    """获取 Doris 客户端实例"""
    return doris_client 